import os
import numpy as np
from mindspore import export, Tensor
from load_models.PNet import PNet
from load_models.RNet import RNet
from load_models.ONet import ONet
import argparse
from mindspore import Tensor, load_checkpoint, load_param_into_net


parser = argparse.ArgumentParser()
parser.add_argument('--model_path', type=str, default='infer_models', help='PNet、RNet、ONet三个模型文件存在的文件夹路径')
args = parser.parse_args()

# 获取P模型
param_pnet = load_checkpoint(os.path.join(args.model_path, 'PNet.ckpt'))
pnet = PNet()
load_param_into_net(pnet, param_pnet)

# 获取R模型
param_rnet = load_checkpoint(os.path.join(args.model_path, 'RNet.ckpt'))
rnet = RNet()
load_param_into_net(rnet, param_rnet)

# 获取O模型
param_onet = load_checkpoint(os.path.join(args.model_path, 'ONet.ckpt'))
onet = ONet()
load_param_into_net(onet, param_onet)

# 将模型由ckpt格式导出为MINDIR格式
input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 12, 12]).astype(np.float32)
export(pnet, Tensor(input_np), file_name="pnet", file_format="MINDIR")

input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 24, 24]).astype(np.float32)
export(rnet, Tensor(input_np), file_name="rnet", file_format="MINDIR")

input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 48, 48]).astype(np.float32)
export(onet, Tensor(input_np), file_name="onet", file_format="MINDIR")